{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Footstep Planning via Mixed-Integer Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import HTML, display\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from matplotlib.patches import Rectangle\n",
    "\n",
    "# drake imports\n",
    "from pydrake.all import MathematicalProgram, OsqpSolver, eq, le, ge\n",
    "from pydrake.solvers import branch_and_bound\n",
    "\n",
    "# increase default size matplotlib figures\n",
    "from matplotlib import rcParams\n",
    "rcParams['figure.figsize'] = (10, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "In this notebook we will implement a simplified footstep planner for a humanoid robot: we will use the method proposed [in this paper](https://groups.csail.mit.edu/robotics-center/public_papers/Deits14a.pdf).\n",
    "The idea is straightforward: we need to plan where to place the feet of the robot in order to move from point A to point B.\n",
    "In doing so, we are allowed to place the feet only in certain safe areas (\"stepping stones\") and each step cannot exceed a certain length.\n",
    "To solve this problem, we will use Mixed-Integer Quadratic Programming (MIQP).\n",
    "\n",
    "MIQP is a relatively nice class of optimization problems.\n",
    "The [branch and bound algorithm](https://en.wikipedia.org/wiki/Branch_and_bound) allows to solve these problems to global optimality, whenever a solution exists, and it certifies infeasibility otherwise.\n",
    "The drawback, however, is that computation times scale exponentially with the number of integer variables in the problem.\n",
    "\n",
    "You will be asked to code most of the components of this MIQP:\n",
    "- The constraint that limit the maximum step length.\n",
    "- The constraint for which a foot cannot be in two stepping stones at the same time.\n",
    "- The constraint that assign each foot to a stepping stone, for each step of the robot.\n",
    "- The objective function that minimizes the sum of the squares of the step lengths.\n",
    "\n",
    "Before moving on, take a look at the following videos to see the Atlas robot using this algorithm!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import IFrame, display\n",
    "display(IFrame(src='https://www.youtube.com/embed/hGhCTPQuMy4', width='560', height='315'))\n",
    "display(IFrame(src='https://www.youtube.com/embed/_6WQxXH-bB4', width='560', height='315'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Terrain\n",
    "\n",
    "We start by constructing the terrain in which the robot will walk.\n",
    "For simplicity, we let the stepping stones be rectangles in the plane.\n",
    "\n",
    "We define each stepping stone by its `center` (2d vector), its `width` (float), and its `height` (float), but we also store [its halfspace representation](https://en.wikipedia.org/wiki/Convex_polytope#Intersection_of_half-spaces).\n",
    "In this representation, a stepping stone is described by a matrix $A$ and a vector $b$ such that a point $x \\in \\mathbb R^2$ lies inside the stepping stone iff $A x \\leq b$.\n",
    "Each row of the matrix $A$ represents one of the four halfspaces that delimit a 2d rectangle.\n",
    "We will need these matrices later in the notebook when we will use an MIP technique known as [the big-M method](https://optimization.mccormick.northwestern.edu/index.php/Disjunctive_inequalities)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SteppingStone(object):\n",
    "\n",
    "    def __init__(self, center, width, height, name=None):\n",
    "        \n",
    "        # store arguments\n",
    "        self.center = center\n",
    "        self.width = width\n",
    "        self.height = height\n",
    "        self.name = name\n",
    "        \n",
    "        # distance from center to corners\n",
    "        c2tr = np.array([width, height]) / 2\n",
    "        c2br = np.array([width, - height]) / 2\n",
    "        \n",
    "        # position of the corners\n",
    "        self.top_right = center + c2tr\n",
    "        self.bottom_right = center + c2br\n",
    "        self.top_left = center - c2br\n",
    "        self.bottom_left = center - c2tr\n",
    "        \n",
    "        # halfspace representation of the stepping stone\n",
    "        self.A = np.array([[1, 0], [0, 1], [-1, 0], [0, -1]])\n",
    "        self.b = np.concatenate([c2tr] * 2) + self.A.dot(center)\n",
    "        \n",
    "    def plot(self, **kwargs):\n",
    "        return plot_rectangle(self.center, self.width, self.height, **kwargs)\n",
    "    \n",
    "# helper function that plots a rectangle with given center, width, and height\n",
    "def plot_rectangle(center, width, height, ax=None, frame=.1, **kwargs):\n",
    "        \n",
    "        # make black the default edgecolor\n",
    "        if not 'edgecolor' in kwargs:\n",
    "            kwargs['edgecolor'] = 'black'\n",
    "            \n",
    "        # make transparent the default facecolor\n",
    "        if not 'facecolor' in kwargs:\n",
    "            kwargs['facecolor'] = 'none'\n",
    "            \n",
    "        # get current plot axis if one is not given\n",
    "        if ax is None:\n",
    "            ax = plt.gca()\n",
    "            \n",
    "        # get corners\n",
    "        c2c = np.array([width, height]) / 2\n",
    "        bottom_left = center - c2c\n",
    "        top_right = center + c2c\n",
    "        \n",
    "        # plot rectangle\n",
    "        rect = Rectangle(bottom_left, width, height, **kwargs)\n",
    "        ax.add_patch(rect)\n",
    "        \n",
    "        # scatter fake corners to update plot limits (bad looking but compact)\n",
    "        ax.scatter(*bottom_left, s=0)\n",
    "        ax.scatter(*top_right, s=0)\n",
    "        \n",
    "        # make axis scaling equal\n",
    "        ax.set_aspect('equal')\n",
    "        \n",
    "        return rect"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the building block for the construction of the robot's terrain, we place the stepping stones.\n",
    "The idea is to replicate the task that Atlas performs in the first video above (at time 1:24).\n",
    "\n",
    "The following class that takes a list of boolean values (e.g. `bool_bridge = [0, 1, 1, 0, 0, 1]`) and generates a collection of stepping stones.\n",
    "We have the `initial` stepping stone on the left, the `goal` stepping stone on the right, the `lateral` stepping stone at the top, and a set of `bridge` stepping stones that connect the `initial` stone to the `goal`.\n",
    "With all the `bridge` stepping stones in place, there would be an easy path for the robot to reach the `goal`.\n",
    "However, out of the potential `len(bool_bridge)` stepping stones forming the bridge, only the ones with entry equal to `1` are actually there.\n",
    "\n",
    "If this description is not super clear, quickly run the next couple of cells and play with the list of booleans in the line `Terrain([1, 0, 1, 1, 0, 1]).plot()`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Terrain(object):\n",
    "    \n",
    "    # parametric construction of the stepping stones\n",
    "    # the following code adapts the position of each stepping\n",
    "    # stone depending on the size and the sparsity of bool_bridge\n",
    "    def __init__(self, bool_bridge):\n",
    "        \n",
    "        # ensure that bool_bridge has only boolean entries\n",
    "        if any(i != bool(i) for i in bool_bridge):\n",
    "            raise ValueError('Entry bool_bridge must be a list of boolean value.')\n",
    "            \n",
    "        # initialize internal list of stepping stones\n",
    "        self.stepping_stones = []\n",
    "        \n",
    "        # add initial stepping stone to the terrain\n",
    "        initial = self.add_stone([0, 0], 1, 1, 'initial')\n",
    "        \n",
    "        # add bridge stepping stones to the terrain\n",
    "        # gap between bridge stones equals bridge stone width\n",
    "        width_bridge = .2\n",
    "        center = initial.bottom_right + np.array([width_bridge * 1.5, initial.height / 4])\n",
    "        centers = [center + np.array([i * 2 * width_bridge, 0]) for i in np.where(bool_bridge)[0]]\n",
    "        self.add_stones(\n",
    "            centers,\n",
    "            [width_bridge] * sum(bool_bridge),\n",
    "            [initial.height / 2] * sum(bool_bridge),\n",
    "            'bridge'\n",
    "        )\n",
    "                \n",
    "        # add goal stepping stone to the terrain\n",
    "        # same dimensions of the initial one\n",
    "        center = initial.center + np.array([initial.width + (len(bool_bridge) * 2 + 1) * width_bridge, 0])\n",
    "        goal = self.add_stone(center, initial.width, initial.height, 'goal')\n",
    "        \n",
    "        # add lateral stepping stone to the terrain\n",
    "        height = .4\n",
    "        clearance = .1\n",
    "        c2g = goal.center - initial.center\n",
    "        width = initial.width + c2g[0]\n",
    "        center = initial.center + c2g / 2 + np.array([0, (initial.height + height) / 2 + clearance])\n",
    "        self.add_stone(center, width, height, 'lateral')\n",
    "        \n",
    "    # adds a stone to the internal list stepping_stones\n",
    "    def add_stone(self, center, width, height, name=None):\n",
    "        stone = SteppingStone(center, width, height, name=name)\n",
    "        self.stepping_stones.append(stone)\n",
    "        return stone\n",
    "    \n",
    "    # adds multiple stones to the internal list stepping_stones\n",
    "    def add_stones(self, centers, widths, heights, name=None):\n",
    "        \n",
    "        # ensure that inputs have coherent size\n",
    "        n_stones = len(centers)\n",
    "        if n_stones != len(widths) or n_stones != len(heights):\n",
    "            raise ValueError('Arguments have incoherent size.')\n",
    "            \n",
    "        # add one stone per time\n",
    "        stones = []\n",
    "        for i in range(n_stones):\n",
    "            stone_name = name if name is None else name + '_' + str(i)\n",
    "            stones.append(self.add_stone(centers[i], widths[i], heights[i], name=stone_name))\n",
    "            \n",
    "        return stones\n",
    "\n",
    "    # returns the stone with the given name\n",
    "    # raise a ValueError if no stone has the given name\n",
    "    def get_stone_by_name(self, name):\n",
    "\n",
    "        # loop through the stones\n",
    "        # select the first with the given name\n",
    "        for stone in self.stepping_stones:\n",
    "            if stone.name == name:\n",
    "                return stone\n",
    "        \n",
    "        # raise error if there is no stone with the given name\n",
    "        raise ValueError(f'No stone in the terrain has name {name}.')\n",
    "    \n",
    "    # plots all the stones in the terrain\n",
    "    def plot(self, title=None, **kwargs):\n",
    "        \n",
    "        # make light green the default facecolor\n",
    "        if not 'facecolor' in kwargs:\n",
    "            kwargs['facecolor'] = [0, 1, 0, .1]\n",
    "        \n",
    "        # plot stepping stones disposition\n",
    "        labels = ['Stepping stone', None]\n",
    "        for i, stone in enumerate(self.stepping_stones):\n",
    "            stone.plot(label=labels[min(i, 1)], **kwargs)\n",
    "            \n",
    "        # set title\n",
    "        plt.title(title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the next cell to play with the list of booleans and make the stones in the bridge appear and disappear.\n",
    "You can also modify the length of the list and the position of the stepping stones will adapt automatically.\n",
    "\n",
    "At the end of the notebook, we will focus on two specific setups: `bool_bridge = [1, 1, 1, 1]` and `bool_bridge = [1, 1, 1, 0]`.\n",
    "In the first case, we expect the robot to walk straight through the bridge to arrive at the goal.\n",
    "In the second, given the strict limits we will enforce on the maximum step length, the robot will have to use the lateral stepping stone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Terrain([1, 0, 1, 1, 0, 1]).plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Mixed-Integer Optimization Problem\n",
    "\n",
    "It's time to write the core of the footstep planner: the MIQP.\n",
    "We start by defining the decision variables.\n",
    "We do this inside of a function, since we ultimately want to define a `footstep_planner` function that given a `terrain` (and a couple of additional parameters) returns the optimal plan to walk through it.\n",
    "\n",
    "Here is the meaning of the variables in the following function:\n",
    "- `n_steps` is the maximum number of steps that the robot can take to reach the goal.\n",
    "- The 2d array `position_left[t]` contains the Cartesian coordinates of the left foot at step `t`.\n",
    "Similarly for `position_right[t]`.\n",
    "- If the binary variable `stone_left[t, i]` is one then, at step `t`, the left foot must be placed on stone `i`.\n",
    "Since a foot cannot be in two stepping stones at the same time, below we will also add the constraint `sum(stone_left[t]) == 1`.\n",
    "Similarly for `stone_right[t, i]`.\n",
    "- The binary `first_left` is one if the first step is taken with the left foot, zero if it is taken with the right foot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_decision_variables(prog, terrain, n_steps):\n",
    "    \n",
    "    # number of stones in the terrain\n",
    "    n_stones = len(terrain.stepping_stones)\n",
    "    \n",
    "    # position of each foot at each step\n",
    "    position_left = prog.NewContinuousVariables(rows=n_steps+1, cols=2)\n",
    "    position_right = prog.NewContinuousVariables(rows=n_steps+1, cols=2)\n",
    "\n",
    "    # binaries that assign feet to stones for each step\n",
    "    stone_left = prog.NewBinaryVariables(rows=n_steps+1, cols=n_stones)\n",
    "    stone_right = prog.NewBinaryVariables(rows=n_steps+1, cols=n_stones)\n",
    "\n",
    "    # which foot to move first\n",
    "    first_left = prog.NewBinaryVariables(1)[0]\n",
    "    \n",
    "    return position_left, position_right, stone_left, stone_right, first_left"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We let the initial position of the feet be close to the center of the `initial` stepping stone.\n",
    "The same for the goal.\n",
    "Then we require that the first and last entries in the `position_left` and `position_right` variables defined above agree with these values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_initial_and_goal_position(prog, terrain, decision_variables):\n",
    "    \n",
    "    # unpack only decision variables needed in this function\n",
    "    position_left, position_right = decision_variables[:2]\n",
    "    \n",
    "    # initial position of the feet of the robot\n",
    "    foot_offset = np.array([0, .2])\n",
    "    center = terrain.get_stone_by_name('initial').center\n",
    "    initial_position_left = center\n",
    "    initial_position_right = center - foot_offset\n",
    "\n",
    "    # goal position of the feet of the robot\n",
    "    center = terrain.get_stone_by_name('goal').center\n",
    "    goal_position_left = center\n",
    "    goal_position_right = center - foot_offset\n",
    "\n",
    "    # enforce initial position of the feet\n",
    "    prog.AddLinearConstraint(eq(position_left[0], initial_position_left))\n",
    "    prog.AddLinearConstraint(eq(position_right[0], initial_position_right))\n",
    "\n",
    "    # enforce goal position of the feet\n",
    "    prog.AddLinearConstraint(eq(position_left[-1], goal_position_left))\n",
    "    prog.AddLinearConstraint(eq(position_right[-1], goal_position_right))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the robot has kinematic limits, its steps cannot be too long.\n",
    "More specifically, at each step `t` from `0` to `n_steps`, we want each foot to lie in a square centered at the other foot.\n",
    "The float `step_span` represents the side of this square (the whole side, not half of it).\n",
    "\n",
    "Your turn to code: in the following cell complete the function `relative_position_limits` to enforce the constraint we just described.\n",
    "So far we have only unpacked the decision variables you need to constrain.\n",
    "\n",
    "P.s.: You might need the constructor `prog.AddLinearConstraint(le(x, y))` to enforce the constraint that the array `x` must be `l`ower than or `e`qual to the array `y`.\n",
    "`ge` can be used similarly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relative_position_limits(prog, n_steps, step_span, decision_variables):\n",
    "    \n",
    "    # unpack only decision variables needed in this function\n",
    "    position_left, position_right = decision_variables[:2]\n",
    "    \n",
    "    # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following constraint is somewhat involved to implement, but the idea behind it is straightforward.\n",
    "If we take the first step with, e.g., the left foot (`first_left == 1`), then we implicitly associated each step in the plan to a foot; using zero-based indexing: 0th step left, 1st step right, 2nd step left, 3rd step right, etc.\n",
    "The following function ensures that, if we move the left foot first, then the left foot will not move during an \"odd\" step and the right foot will not move during an \"even\" step.\n",
    "Similarly if we decide to take the first step with the right foot.\n",
    "\n",
    "Here is how we implemented it.\n",
    "At each step `t`, we require the next position of a foot to lie in a square centered at its current position.\n",
    "We then use the binaries `first_left` and `first_right = 1 - first_left` to make the side of this square sufficiently large or zero, depending on whether the foot is allowed to move at step `t`.\n",
    "Take a couple of minutes to double check that we got the implementation right!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step_sequence(prog, n_steps, step_span, decision_variables):\n",
    "    \n",
    "    # unpack only decision variables needed in this function\n",
    "    position_left, position_right = decision_variables[:2]\n",
    "    first_left = decision_variables[-1]\n",
    "    \n",
    "    # variable equal to one if first step is with right foot\n",
    "    first_right = 1 - first_left\n",
    "    \n",
    "    # note that the step_span coincides with the maximum distance\n",
    "    # (both horizontal and vertical) between the position of\n",
    "    # a foot at step t and at step t + 1\n",
    "    step_limit = np.ones(2) * step_span\n",
    "\n",
    "    # sequence for the robot steps implied by the binaries first_left and first_right\n",
    "    # (could be written more compactly, but would be harder to read)\n",
    "    for t in range(n_steps):\n",
    "\n",
    "        # lengths of the steps\n",
    "        step_left = position_left[t + 1] - position_left[t]\n",
    "        step_right = position_right[t + 1] - position_right[t]\n",
    "\n",
    "        # for all even steps\n",
    "        if t % 2 == 0:\n",
    "            limit_left = step_limit * first_left   # left foot can move iff first_left\n",
    "            limit_right = step_limit * first_right # right foot can move iff first_right\n",
    "\n",
    "        # for all odd steps\n",
    "        else:\n",
    "            limit_left = step_limit * first_right # left foot can move iff first_right\n",
    "            limit_right = step_limit * first_left # right foot can move iff first_left\n",
    "\n",
    "        # constraints on left-foot relative position\n",
    "        prog.AddLinearConstraint(le(step_left, limit_left))\n",
    "        prog.AddLinearConstraint(ge(step_left, - limit_left))\n",
    "        \n",
    "        # constraints on right-foot relative position\n",
    "        prog.AddLinearConstraint(le(step_right, limit_right))\n",
    "        prog.AddLinearConstraint(ge(step_right, - limit_right))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's your turn to code again.\n",
    "As mentioned before, a foot cannot be in two stepping stones at the same time.\n",
    "This can be enforced by requiring that, for each step `t`, the binaries `stone_left[t]` sum up to one.\n",
    "The same for `stone_right[t]`.\n",
    "Implement this linear constraint in the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_stone_per_foot(prog, n_steps, decision_variables):\n",
    "    \n",
    "    # unpack only decision variables needed in this function\n",
    "    stone_left, stone_right = decision_variables[2:4]\n",
    "    \n",
    "    # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One more constraint to add: \"if the binary `stone_left[t, i]` is one, then the left foot at step `t` must lie on the `i`th stepping stone\".\n",
    "To enforce this constraint we use the [the big-M method](https://optimization.mccormick.northwestern.edu/index.php/Disjunctive_inequalities).\n",
    "\n",
    "Think of each stepping stone $S_i$ as a subset of $\\mathbb R^2$.\n",
    "Using the halfspace representation,\n",
    "$$\n",
    "S_i = \\{ x \\in \\mathbb R^2 : A_i x \\leq b_i \\}.\n",
    "$$\n",
    "Let $x_t$ be the position of, e.g., the left foot at step $t$.\n",
    "The constraint $x_t \\in \\bigcup_i S_i$ can be enforced as follows.\n",
    "We define a binary $\\delta_{t, i} \\in \\{0, 1\\}$ per stepping stone (these are the binaries `stone_left[t, i]` we defined above).\n",
    "We add the constraint\n",
    "$$\n",
    "A_i x_t \\leq b_i + (1 - \\delta_{t, i}) M_i\n",
    "$$\n",
    "for all $t$ and $i$, where $M_i$ is a vector of \"sufficiently large\" positive constants.\n",
    "Additionally, as we already did in the last cell, we require\n",
    "$$\n",
    "\\sum_i \\delta_{t, i} = 1\n",
    "$$\n",
    "for all $t$.\n",
    "\n",
    "Do you see how these two constraints do the job?\n",
    "If $\\delta_{t, i} = 1$, from the last equation, we must have $\\delta_{t, j} = 0$ for all $j \\neq i$.\n",
    "This implies\n",
    "$$\n",
    "A_i x_t \\leq b_i, \\quad\n",
    "A_j x_t \\leq b_j + M_j.\n",
    "$$\n",
    "Therefore $x_t$ belongs to $S_i$, while, since $M_j$ is a vector of large constants, the second inequality is redundant and does not constrain the value of $x_t$.\n",
    "\n",
    "It's very important to notice that these constraints are linear in our decision variables ($x_t$ and $\\delta_{t, i}$).\n",
    "Mixed-integer programs with nonlinear constraints are extremely hard problems, especially if some of the constraints are nonconvex.\n",
    "\n",
    "One issue left: what's the magic value for the constants $M_i$?\n",
    "One would be tempted to put a huge number there, but unfortunately the larger the $M_i$ the harder the problem is to be solved via branch and bound.\n",
    "In the following cell we wrote a function that returns a vector $M$ that can be used for all $i$.\n",
    "(This is possible since our stepping stones are rectangles and they all have the same number of facets, 4, which is equal to the number of rows in $A_i$).\n",
    "If you are not a MIP enthusiast, do not feel the need to understand the logic behind the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter for the big-M method\n",
    "# carefully chosen for the terrain above\n",
    "def get_big_M(terrain):\n",
    "    \n",
    "    # big-M parameter for the horizontal axis\n",
    "    initial = terrain.get_stone_by_name('initial')\n",
    "    goal = terrain.get_stone_by_name('goal')\n",
    "    M = [goal.center[0] - initial.center[0]]\n",
    "    \n",
    "    # big-M parameter for the vertical axis\n",
    "    lateral = terrain.get_stone_by_name('lateral')\n",
    "    M.append(lateral.top_right[1] - initial.center[1])\n",
    "    \n",
    "    return np.array(M * 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell, for each step $t$ and stone $i$, implement the big-M constraint\n",
    "$$\n",
    "A_i x_t \\leq b_i + (1 - \\delta_{t, i}) M,\n",
    "$$\n",
    "where $x_t$ represents the variable `position_left[t]` (respectively `position_right[t]`), $\\delta_{t, i}$ is `stone_left[t, i]` (respectively `stone_right[t, i]`), and $M$ is obtained via the function `get_big_M` we just defined.\n",
    "The arrays $A_i$ and $b_i$ are stored as attributes of the class `SteppingStone`.\n",
    "Stones are assumed to be ordered as in the list `terrain.stepping_stones`, meaning that the $i$th stone here is `terrain.stepping_stones[i]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def foot_in_stepping_stone(prog, terrain, n_steps, decision_variables):\n",
    "    \n",
    "    # unpack only decision variables needed in this function\n",
    "    position_left, position_right, stone_left, stone_right = decision_variables[:4]\n",
    "    \n",
    "    # big-M vector\n",
    "    M = get_big_M(terrain)\n",
    "    \n",
    "    # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last piece of code you are asked to write.\n",
    "In order to promote small steps, add a quadratic objective that minimizes the step length.\n",
    "More precisely, for each step `t`, add to the cost function the square of the two norm of `position_left[t + 1] - position_left[t]` and the square of the two norm of `position_right[t + 1] - position_right[t]`.\n",
    "To do this, you can use the `MathematicalProgram` method `AddQuadraticCost`.\n",
    "\n",
    "For the moment we added a dummy objective function.\n",
    "This is just needed to let the notebook run without errors: erase it to write your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def minimize_step_length(prog, n_steps, decision_variables):\n",
    "    \n",
    "    # unpack only decision variables needed in this function\n",
    "    position_left, position_right = decision_variables[:2]\n",
    "    \n",
    "    # dummy objective needed to let the notebook run without errors # erase here\n",
    "    prog.AddQuadraticCost(position_left.flatten().dot(position_left.flatten())) # erase here\n",
    "    prog.AddQuadraticCost(position_right.flatten().dot(position_right.flatten())) # erase here\n",
    "    \n",
    "    # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The footstep planner is ready!\n",
    "In the following cell we just put together all the pieces we wrote so far and solve the optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def footstep_planner(terrain, n_steps, step_span):\n",
    "    \n",
    "    # initialize optimization problem\n",
    "    prog = MathematicalProgram()\n",
    "    \n",
    "    # optimization variables\n",
    "    decision_variables = add_decision_variables(prog, terrain, n_steps)\n",
    "    \n",
    "    # constraints\n",
    "    set_initial_and_goal_position(prog, terrain, decision_variables)\n",
    "    relative_position_limits(prog, n_steps, step_span, decision_variables)\n",
    "    step_sequence(prog, n_steps, step_span, decision_variables)\n",
    "    one_stone_per_foot(prog, n_steps, decision_variables)\n",
    "    foot_in_stepping_stone(prog, terrain, n_steps, decision_variables)\n",
    "    \n",
    "    # objective function\n",
    "    minimize_step_length(prog, n_steps, decision_variables)\n",
    "    \n",
    "    # solve\n",
    "    bb = branch_and_bound.MixedIntegerBranchAndBound(prog, OsqpSolver().solver_id())\n",
    "    result = bb.Solve()\n",
    "    \n",
    "    # ensure that the problem is feasible\n",
    "    if result != result.kSolutionFound:\n",
    "        raise ValueError('Infeasible optimization problem.')\n",
    "\n",
    "    # retrieve result of the optimization\n",
    "    decision_variables_opt = [bb.GetSolution(v) for v in decision_variables]\n",
    "    objective_opt = bb.GetOptimalCost()\n",
    "    \n",
    "    return decision_variables_opt, objective_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try the Footstep Planner on two Terrains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We consider the following two terrains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# complete bridge\n",
    "terrain_A = Terrain([1, 1, 1, 1])\n",
    "terrain_A.plot('Terrain A')\n",
    "plt.show()\n",
    "\n",
    "# one stepping stone missing in the bridge\n",
    "terrain_B = Terrain([1, 1, 1, 0])\n",
    "terrain_B.plot('Terrain B')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the additional parameters for the footstep planner.\n",
    "`n_steps` is chosen to make the optimizations solvable in a few seconds, `step_span` to achieve a specific behaviour you'll see in a couple of cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# maximum number of steps to reach the goal\n",
    "n_steps = 8\n",
    "\n",
    "# side of the square that limits each step\n",
    "step_span = .8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an animation function that you can use to check your result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def animate_footstep_plan(terrain, step_span, position_left, position_right, title=None):\n",
    "\n",
    "    # initialize figure for animation\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    # plot stepping stones\n",
    "    terrain.plot(title=title, ax=ax)\n",
    "\n",
    "    # initial position of the feet\n",
    "    left_foot = ax.scatter(0, 0, color='r', zorder=3, label='Left foot')\n",
    "    right_foot = ax.scatter(0, 0, color='b', zorder=3, label='Right foot')\n",
    "\n",
    "    # initial step limits\n",
    "    left_limits = plot_rectangle(\n",
    "        [0 ,0],    # center\n",
    "        step_span, # width\n",
    "        step_span, # eight\n",
    "        ax=ax,\n",
    "        edgecolor='b',\n",
    "        label='Left-foot limits'\n",
    "    )\n",
    "    right_limits = plot_rectangle(\n",
    "        [0 ,0],    # center\n",
    "        step_span, # width\n",
    "        step_span, # eight\n",
    "        ax=ax,\n",
    "        edgecolor='r',\n",
    "        label='Right-foot limits'\n",
    "    )\n",
    "\n",
    "    # misc settings\n",
    "    plt.close()\n",
    "    ax.legend(loc='upper left', bbox_to_anchor=(0, 1.3), ncol=2)\n",
    "\n",
    "    def animate(n_steps):\n",
    "\n",
    "        # scatter feet\n",
    "        left_foot.set_offsets(position_left[n_steps])\n",
    "        right_foot.set_offsets(position_right[n_steps])\n",
    "\n",
    "        # limits of reachable set for each foot\n",
    "        c2c = np.ones(2) * step_span / 2\n",
    "        right_limits.set_xy(position_left[n_steps] - c2c)\n",
    "        left_limits.set_xy(position_right[n_steps] - c2c)\n",
    "\n",
    "    # create ad display animation\n",
    "    ani = FuncAnimation(fig, animate, frames=n_steps+1, interval=1e3)\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "    \n",
    "def generate_and_animate_footstep_plan(terrain, n_steps, step_span, title=None):\n",
    "    \n",
    "    # run footstep planner\n",
    "    decision_variables, objective = footstep_planner(terrain, n_steps, step_span)\n",
    "    \n",
    "    # animate result\n",
    "    animate_footstep_plan(terrain, step_span, *decision_variables[:2], title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the animation of the footstep planner for the two benchmark terrains.\n",
    "If you did things correctly, you should see the robot crossing the bridge for `terrain_A`.\n",
    "Whereas the max step limit requires it to take the less efficient route at the top of the terrain (`lateral` stone) in case of `terrain_B`.\n",
    "Note that a local solver (as opposed to branch and bound which is a global optimizer) would have had a very hard time realizing this!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_and_animate_footstep_plan(terrain_A, n_steps, step_span, 'Terrain A')\n",
    "generate_and_animate_footstep_plan(terrain_B, n_steps, step_span, 'Terrain B')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final Comments\n",
    "\n",
    "In this notebook we used a \"simple\" branch-and-bound solver (written in Drake by Hongkai Dai).\n",
    "If you want to solve this problem for larger terrains, you might need something more advanced.\n",
    "Drake supports two state-of-the-art solvers for mixed-integer programming (both free for academic use): [Gurobi](https://www.gurobi.com) and [Mosek](https://www.mosek.com).\n",
    "With these, you should be able to approximately double the stepping stones in these terrains and the maximum number of steps.\n",
    "\n",
    "Here we used the big-M method, which is the simplest way in which we can transcribe the the footstep planning problem into an MIQP.\n",
    "Much more efficient transcription techniques are available.\n",
    "If you are interested, have a look at [this very nice survey paper](https://dspace.mit.edu/bitstream/handle/1721.1/96480/Vielma-2015-Mixed%20Integer%20Linear.pdf?sequence=1&isAllowed=y)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "\n",
    "You can check your work by running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.humanoids.footstep_planning.test_footstep_planning import TestFootstepPlanning\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestFootstepPlanning], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}